#!/usr/bin/env python
# coding: utf-8

# 数值方法3：偏微分方程1 使用有限差分法解一维热传导（扩散）方程
# 有热源情况

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
from matplotlib.animation import FFMpegWriter

# 热传导参数
a = [0.000000198499152747519, 0.000000204397304165286, 0.000000351372539220984, 0.0000236107597605194]
dx = 0.0001
X = [0.0006, 0.006, 0.0036, 0.005]
x = []
for item in range(len(X)):
    if item == 0:
        x.append(np.linspace(0, X[item], int(X[item] / dx + 1)))
    else:
        x.append(np.linspace(X[item - 1], X[item], int(X[item] / dx + 1)))
# x = np.linspace(0, X, int(X / dx + 1))
dt = 1
T = 5400
t = np.linspace(0, T, int(T / dt + 1))
# print(x)
# 算子
A = (-2 * np.eye(len(x[0]), len(x[0]), dtype=int)
     + np.diag(np.diag(np.ones([len(x[0]), len(x[0])]), k=1), k=1)
     + np.diag(np.diag(np.ones([len(x[0]), len(x[0])]), k=1), k=-1))
# print(A)
# 初始化
u = 0 * x
# 边界
m1 = 0+0.0*np.sin(t)  # 热流流入
m2 = 0-0.0*np.sin(8*t)

# 热源
list_f = [75 for i in range(len(x[0]))]
f = np.mat(list_f).T
# f = 10*np.mat(5*np.exp(-20*(np.power(x-0.5, 2)))).T
# 解方程组
list_u = [37 for i in range(len(x[0]))]
u = np.mat(list_u).T
# print(u.shape)
for n in range(len(t) - 1):
    idx = 0
    if n>int((len(t)-1)/4) & n<int((len(t)-1)/2):
        idx = 1
    elif n>int((len(t)-1)/2) & n<int(3*(len(t)-1)/4):
        idx = 2
    else:
        idx = 3
    u = np.append(u, np.add(u[:, n], np.add(a[idx] * a[idx] / dx / dx * np.mat(A) * u[:, n], f) * dt), axis=1)
    # u = np.append(u, u[:, n] + a[idx] * a[idx] * dt / dx / dx * np.mat(A) * u[:, n], axis=1)
    u[0, n + 1] = m1[n + 1]
    u[-1, n + 1] = m2[n + 1]
    print(u)
# print(u)
# 绘制边界线

Y = np.array(u)
X = np.array(np.linspace(0, 0.0152, Y.shape[0]))
plt.plot(X, Y[:, -1])
plt.xlabel("houdu")
plt.ylabel("wendu")
plt.title('FinalStates Temperature')
plt.savefig("FinalStatesT.png")  # 保存图片
plt.show()

# 3D曲面
# get_ipython().run_line_magic('matplotlib', 'notebook')
# 准备数据
x_x, y_y = np.meshgrid(t, X)

# 绘制图片
fig = plt.figure("3D Surface", facecolor="lightgray")
plt.title("Temperature3D", fontsize=18)

# 设置为3D图片类型
ax3d = Axes3D(fig, auto_add_to_figure=False)
fig.add_axes(ax3d)

ax3d.set_xlabel("time")
ax3d.set_ylabel("x")
ax3d.set_zlabel("temperature")
plt.tick_params(labelsize=10)

ax3d.plot_surface(x_x, y_y, u, cstride=20, rstride=20, cmap="jet")

plt.savefig("Temperature3D.png")  # 保存图片
plt.show()
'''
# 初始化画布
fig = plt.figure()
# plt.xlim(0,X)
plt.ylim(0, np.max(u))  # np.max(psi)
plt.grid(ls='--')

ys = u[:, 0]
Figure = plt.plot(x, ys, c='blue', alpha=0.8)[0]


# 更新函数
def updata(num):
    ys = u[:, num]
    Figure.set_data(x, ys)
    return Figure


# animation库绘制动图
ani = animation.FuncAnimation(fig=fig, func=updata, frames=np.arange(0, int(len(t) / 2)),
                              interval=2)  # frames帧数，interval间隔（帧）
plt.show()

## 保存为mp4，运行速度较慢，不保存时注释掉。
# mywriter = FFMpegWriter(fps=60)
# ani.save('TemperatureHistory.MP4',writer=mywriter)


'''